import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import time

# ====================== 修复1：中文显示问题 ======================
# 设置支持中文的字体（添加更多备选字体）
plt.rcParams['font.sans-serif'] = ['Microsoft YaHei', 'SimHei', 'KaiTi', 'FangSong', 'Arial Unicode MS']
plt.rcParams['axes.unicode_minus'] = False

class Drone:
    """无人机物理模型"""
    def __init__(self):
        self.roll_angle = 0.0  # 横滚角度 (弧度)
        self.roll_rate = 0.0   # 横滚角速度 (弧度/秒)
        self.Ixx = 0.1         # 转动惯量
        self.damping = 0.05    # 阻尼系数
        
    def update(self, torque, dt):
        """根据施加的力矩更新无人机状态"""
        # 角加速度 = (扭矩 - 阻尼 * 角速度) / 转动惯量
        roll_acc = (torque - self.damping * self.roll_rate) / self.Ixx
        
        # 更新角速度和角度
        self.roll_rate += roll_acc * dt
        self.roll_angle += self.roll_rate * dt
        
        return self.roll_angle, self.roll_rate

class PIDController:
    """PID控制器"""
    def __init__(self, Kp, Ki, Kd, limit):
        self.Kp = Kp
        self.Ki = Ki
        self.Kd = Kd
        self.limit = limit  # 输出限幅
        self.integral = 0
        self.prev_error = 0
        
    def compute(self, setpoint, measurement, dt):
        """计算控制输出"""
        error = setpoint - measurement
        
        # 比例项
        P = self.Kp * error
        
        # 积分项
        self.integral += error * dt
        I = self.Ki * self.integral
        
        # 微分项
        derivative = (error - self.prev_error) / dt if dt > 0 else 0
        D = self.Kd * derivative
        self.prev_error = error
        
        # 总和并限幅
        output = P + I + D
        output = np.clip(output, -self.limit, self.limit)
        
        return output, error, P, I, D

# 创建无人机和控制器实例
drone = Drone()

# 外环PID控制器 (角度环)
outer_pid = PIDController(Kp=2.5, Ki=0.1, Kd=0.8, limit=2.0)

# 内环PID控制器 (角速度环)
inner_pid = PIDController(Kp=0.8, Ki=0.5, Kd=0.05, limit=5.0)

# 仿真参数
dt = 0.02  # 时间步长 (秒)
total_time = 10  # 总仿真时间 (秒)
steps = int(total_time / dt)

# 存储数据用于绘图
time_points = np.zeros(steps)
setpoint_history = np.zeros(steps)
angle_history = np.zeros(steps)
rate_history = np.zeros(steps)
torque_history = np.zeros(steps)
outer_output_history = np.zeros(steps)
error_history = np.zeros(steps)

# 创建图形
fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(12, 8))
fig.suptitle('无人机串级PID控制仿真', fontsize=16)

# 初始化绘图
ax1.set_xlim(0, total_time)
ax1.set_ylim(-0.5, 1.5)
ax1.set_ylabel('角度 (rad)')
ax1.set_title('角度响应')
ax1.grid(True)

ax2.set_xlim(0, total_time)
ax2.set_ylim(-3, 3)
ax2.set_ylabel('角速度 (rad/s)')
ax2.set_title('角速度响应')
ax2.grid(True)

ax3.set_xlim(0, total_time)
ax3.set_ylim(-6, 6)
ax3.set_xlabel('时间 (s)')
ax3.set_ylabel('扭矩 (Nm)')
ax3.set_title('控制输出')
ax3.grid(True)

ax4.set_xlim(0, total_time)
ax4.set_ylim(-1, 1)
ax4.set_xlabel('时间 (s)')
ax4.set_ylabel('误差 (rad)')
ax4.set_title('角度误差')
ax4.grid(True)

# 初始化线条
line1, = ax1.plot([], [], 'b-', label='实际角度')
line1_sp, = ax1.plot([], [], 'r--', label='期望角度')
ax1.legend(loc='upper right')

line2, = ax2.plot([], [], 'g-', label='角速度')
line2_sp, = ax2.plot([], [], 'r--', label='角速度设定值')
ax2.legend(loc='upper right')

line3, = ax3.plot([], [], 'm-', label='扭矩')
ax3.legend(loc='upper right')

line4, = ax4.plot([], [], 'r-', label='角度误差')
ax4.legend(loc='upper right')

# 添加文本显示当前PID值
pid_text = ax1.text(0.02, 0.95, '', transform=ax1.transAxes, fontfamily='monospace')

def init():
    """初始化动画"""
    line1.set_data([], [])
    line1_sp.set_data([], [])
    line2.set_data([], [])
    line2_sp.set_data([], [])
    line3.set_data([], [])
    line4.set_data([], [])
    pid_text.set_text('')
    return line1, line1_sp, line2, line2_sp, line3, line4, pid_text

def animate(i):
    """动画更新函数"""
    # 设置点 - 方波信号
    if i < steps//3:
        setpoint = 0.0
    elif i < 2*steps//3:
        setpoint = 1.0
    else:
        setpoint = 0.0
    
    # 串级PID控制
    # 外环：角度环，输出是角速度设定值
    outer_output, error, P_outer, I_outer, D_outer = outer_pid.compute(
        setpoint, drone.roll_angle, dt)
    
    # 内环：角速度环，输出是扭矩
    torque, _, P_inner, I_inner, D_inner = inner_pid.compute(
        outer_output, drone.roll_rate, dt)
    
    # 更新无人机状态
    angle, rate = drone.update(torque, dt)
    
    # 存储数据
    time_points[i] = i * dt
    setpoint_history[i] = setpoint
    angle_history[i] = angle
    rate_history[i] = rate
    torque_history[i] = torque
    outer_output_history[i] = outer_output
    error_history[i] = error
    
    # 更新绘图
    line1.set_data(time_points[:i+1], angle_history[:i+1])
    line1_sp.set_data(time_points[:i+1], setpoint_history[:i+1])
    line2.set_data(time_points[:i+1], rate_history[:i+1])
    line2_sp.set_data(time_points[:i+1], outer_output_history[:i+1])
    line3.set_data(time_points[:i+1], torque_history[:i+1])
    line4.set_data(time_points[:i+1], error_history[:i+1])
    
    # 更新PID值显示
    pid_info = (f"外环: P={P_outer:.2f}, I={I_outer:.2f}, D={D_outer:.2f}\n"
                f"内环: P={P_inner:.2f}, I={I_inner:.2f}, D={D_inner:.2f}")
    pid_text.set_text(pid_info)
    
    return line1, line1_sp, line2, line2_sp, line3, line4, pid_text

# 创建动画
ani = FuncAnimation(fig, animate, frames=steps, init_func=init, 
                    interval=dt*1000, blit=True, repeat=False)

plt.tight_layout()
plt.show()
